import re
from langchain_ollama import ChatOllama
llm = ChatOllama(
  model="deepseek-r1:8b",
  # model="gemma3:1b",
  base_url="http://localhost:11434",  # Ollama 服务地址
  temperature=0.7,  # 创造性程度
  num_predict=1024,  # 最大生成长度
  options={'stop': ['<think>','</think>']}, # 遇到这些标记就停止
)

# 创建一个空的DataFrame用于存储结果
import pandas as pd
df = pd.DataFrame(columns=["flower_type", "price", "description", "reason"])

# 数据准备
# flowers = ["玫瑰", "百合", "康乃馨"]
# prices = ["50", "30", "20"]
flowers = ["玫瑰"]
prices = ["50"]

# 定义我们想要接收的数据格式
from pydantic import BaseModel, Field
class FlowerDescription(BaseModel):
  flower_type: str = Field(description="鲜花的种类")
  price: int = Field(description="鲜花的价格")
  description: str = Field(description="鲜花的描述文案")
  reason: str = Field(description="为什么要这样写这个文案")

# ------Part 3
# 创建输出解析器
from langchain.output_parsers import PydanticOutputParser
output_parser = PydanticOutputParser(pydantic_object=FlowerDescription)

# 获取输出格式指示
format_instructions = output_parser.get_format_instructions()
# 打印提示
# print("输出格式：",format_instructions)

# ------Part 4
# 创建提示模板
from langchain.prompts import PromptTemplate
prompt_template = """您是一位专业的鲜花店文案撰写员。
对于售价为 {price} 元的 {flower} ，您能提供一个吸引人的简短中文描述吗？
{format_instructions}"""

# 根据模板创建提示，同时在提示中加入输出解析器的说明
prompt = PromptTemplate.from_template(prompt_template, 
  partial_variables={"format_instructions": format_instructions}) 


# ------Part 5
for flower, price in zip(flowers, prices):
  # 根据提示准备模型的输入
  input_promp = prompt.format(flower=flower, price=price)
  # 打印提示
  # print("提示：", input_promp)

  # 获取模型的输出
  output = llm.invoke(input_promp)

  # 解析模型的输出
  try:
    parsed_output = output_parser.parse(output.content)
    print('parsed_output',parsed_output)
    # parsed_output_dict = parsed_output.model_dump()  # 将Pydantic格式转换为字典
  except Exception as e:
    print(f"解析输出时出错: {str(e)}")
    continue
    
  df.loc[len(df)] = parsed_output.model_dump()

# 打印字典
print("输出的数据：", df.to_dict(orient='records'))